#include <iostream>
#include <random>
#include <ctime>
#include <sys/time.h>

#include <arm_neon.h>

using namespace std;

#define PERF_CPU(name, count) PerfCPU perf_CPU_##name##__COUNTER__(#name, count)

class PerfCPU
{
public:
    PerfCPU(const std::string& name, int loop) { 
        m_name = name;
        m_loop = loop;
        gettimeofday(&m_start, NULL);
    }

    ~PerfCPU() {
        gettimeofday(&m_end, NULL);
        float elapsed_time = (m_end.tv_sec - m_start.tv_sec) * 1000.0 + (m_end.tv_usec - m_start.tv_usec) / 1000.0;
        std::cout << m_name << " elapse: " << elapsed_time / m_loop << " ms" << std::endl;
    }

private:
    std::string m_name;
    struct timeval m_start, m_end;
    int m_loop = 1;
}; // class PerfCPU

typedef void (*MEMCPY64_FUNC)(float*, const float*);

void memcpy_float32x64_naive(float* dst, const float* src)
{
    // 循环展开
    #pragma unroll
    for (int j = 0; j < 64; ++j) {
        dst[j] = src[j];
    }
}

void memcpy_float32x64_neon(float* dst, const float* src)
{
    // 循环展开
    #pragma unroll
    for (int j = 0; j < 64 / 16; ++j) {
        float32x4x4_t v0 = vld1q_f32_x4(src + j * 16);
        vst1q_f32_x4(dst + j * 16, v0);
    }
}

/**
inline void memcpy_float32x64_asm(float* __restrict__ dst, const float* __restrict__ src)
{
    asm volatile("\r\n"
        "ldp q0,  q1,   [%0]\r\n"
        "stp q0,  q1,   [%1]\r\n"
        
        "ldp q2,  q3,   [%0, 32 ] \r\n"
        "stp q2,  q3,   [%1, 32 ] \r\n"

        "ldp q4,  q5,   [%0, 64 ] \r\n"
        "stp q4,  q5,   [%1, 64 ] \r\n"

        "ldp q6,  q7,   [%0, 96 ] \r\n"
        "stp q6,  q7,   [%1, 96 ] \r\n"

        "ldp q8,  q9,   [%0, 128] \r\n"
        "stp q8,  q9,   [%1, 128] \r\n"

        "ldp q10, q11,  [%0, 160] \r\n"
        "stp q10, q11,  [%1, 160] \r\n"

        "ldp q12, q13,  [%0, 192] \r\n"
        "stp q12, q13,  [%1, 192] \r\n"

        "ldp q14, q15,  [%0, 224] \r\n"        
        "stp q14, q15,  [%1, 224] \r\n"
        : "=r"(src), "=r"(dst)
        : "0"(src), "1"(dst)
        : "memory", "q0", "q1", "q2", "q3", "q4", "q5", "q6", "q7", "q8", "q9", "q10", "q11",
            "q12", "q13", "q14", "q15"
    );
}
*/

void SimpleCopyKernel(int volume_size, float* dst, const float* src)
{
    // #pragma parallel for num_threads(4)
    for (int i = 0; i < volume_size; ++i) {
        const float* src_ptr = src + (i * 64);
        float* dst_ptr = dst + (i * 64);
        memcpy_float32x64_neon(dst_ptr, src_ptr);
    }
}

float* SimpleKernelLaunch(int volume_size, const float* src)
{
    float* output = new float[volume_size * 64];

    std::cout << "volume_size: " << volume_size << std::endl;
    int loop = 10;
    PERF_CPU(single_test, loop);
    for (int i = 0; i < loop; ++i) {
        SimpleCopyKernel(volume_size, output, src);
    }
    return output;
}


float* lut(int volume_size, int32_t* img_idx,
    int32_t* lut_idx, float* feature, MEMCPY64_FUNC func)
{
    float* output_cpu = new float[volume_size * 64];
    memset(output_cpu, 0.0, sizeof(float) * volume_size * 64);
    // calculate cpu

    int loop = 10;
    PERF_CPU(lut_naive, loop);
    for (int l = 0; l < loop; ++l) {
        // 开启四线程
        #pragma omp parallel for num_threads(2)
        for (int i = 0; i < volume_size; ++i) {
            if (img_idx[i] < 0) {
                continue;
            }

            int idx = img_idx[i];
            float* src = feature + idx * 64 * 176 * 64 + lut_idx[i] * 64;
            float* dst = output_cpu + i * 64;

            func(dst, src);
        }
    }
    return output_cpu;
}

int main(int argc, char* argv[])
{
    // nhwc
    float* feature = new float[6 * 64 * 176 * 64];

    std::default_random_engine generator;
    std::uniform_real_distribution<float> distribution(0.0, 1.0);

    for (int i = 0; i < 6 * 64 * 176 * 64; ++i) {
        feature[i] = distribution(generator);
    }
    int volume_size = 200 * 200 * 4;
    int32_t* img_idx = new int32_t[volume_size];
    for (int i = 0; i < volume_size; ++i) {
        img_idx[i] = rand() % 6;
    }
    int32_t* lut_idx = new int32_t[volume_size];
    for (int i = 0; i < volume_size; ++i) {
        lut_idx[i] = rand() % (64 * 176);
    }
    
    float* output1 = lut(volume_size, img_idx, lut_idx, feature, memcpy_float32x64_naive);
    float* output2 = lut(volume_size, img_idx, lut_idx, feature, memcpy_float32x64_neon);

    float* feature2 = new float[volume_size];
    SimpleKernelLaunch(volume_size, feature2);

    // print
    for (int i = 0; i < 1; ++i) {
        for (int c = 0; c < volume_size; ++c) {
            float value = std::abs(output1[i * 64 + c] - output2[i * 64 + c]);
            if (value > 1e-5)
            {
                printf("error at %d, %d, %f, %f\n", i, c, output1[i * 64 + c], output2[i * 64 + c]);
            }
        }
    }

    return 0;
}